import numpy as np
import matplotlib.pyplot as plt
from scipy.stats import hypergeom, binom
import matplotlib
matplotlib.use(backend="TkAgg")

# 参数
N = 52  # 总数
K = 4  # A 的数量
p = K / N

# 三种抽牌数
n_list = [5, 20, 40]

plt.figure(figsize=(12, 8))

for i, n in enumerate(n_list, 1):
    x = np.arange(0, n + 1)

    # 超几何 & 二项分布 PMF(probability mass function)
    pmf_hyper = hypergeom.pmf(x, N, K, n)
    pmf_binom = binom.pmf(x, n, p)

    # 期望
    mean_hyper = hypergeom.mean(N, K, n)
    mean_binom = binom.mean(n, p)

    # 绘图
    plt.subplot(2, 2, i)
    plt.plot(x, pmf_hyper, 'o-', label=f'Hypergeometric (mean={mean_hyper:.2f})')
    plt.plot(x, pmf_binom, 's--', label=f'Binomial (mean={mean_binom:.2f})')

    plt.axvline(mean_hyper, color='blue', linestyle=':', alpha=0.7)
    plt.axvline(mean_binom, color='orange', linestyle=':', alpha=0.7)

    plt.xticks(x)
    plt.xlabel("Number of Aces (X)")
    plt.ylabel("Probability")
    plt.title(f"n = {n} draws")
    plt.legend()
    plt.grid(alpha=0.3)

plt.tight_layout()
plt.show()

'''
图像解释

n=5
超几何分布与二项几乎重叠。
因为抽牌数远小于总数，不放回的影响小。

n=20
两者的差异开始明显，超几何分布更集中，尾部概率更小。
因为抽出的牌数已经占总数相当比例，不放回导致“负相关”效应变强。

n=40
差异非常大：
超几何分布几乎收缩到某个区间（不可能出现太多 A 或太少 A）。
二项分布仍然比较“宽”，允许更极端的情况。
'''